# simple_network_utils.py - 简化的网络工具
"""
简化版网络工具，提高兼容性
专门解决 Hugging Face Hub 上传时的网络问题
"""

import time
import random
import requests
import os
from typing import Callable, Any, Optional, List, Dict
from functools import wraps


class SimpleNetworkRetryConfig:
    """简化的网络重试配置"""
    def __init__(
        self,
        max_retries: int = 5,
        base_delay: float = 2.0,
        max_delay: float = 120.0,
        backoff_factor: float = 2.5,
        jitter: bool = True
    ):
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.max_delay = max_delay
        self.backoff_factor = backoff_factor
        self.jitter = jitter


def calculate_delay(attempt: int, config: SimpleNetworkRetryConfig) -> float:
    """计算重试延迟时间"""
    delay = min(config.base_delay * (config.backoff_factor ** attempt), config.max_delay)
    
    if config.jitter:
        # 添加随机抖动
        delay = delay * (0.5 + random.random() * 0.5)
    
    return delay


def is_network_error(exception: Exception) -> bool:
    """判断是否为网络错误"""
    error_msg = str(exception).lower()
    network_keywords = [
        'connection reset',
        'connection broken', 
        'connection aborted',
        'connection timeout',
        'proxy error',
        'max retries exceeded',
        'timeout',
        'network unreachable',
        'ssl',
        'certificate',
        'unable to connect',
        'connection error',
        'read timed out',
        'connection pool',
        'httpsconnectionpool'
    ]
    
    return any(keyword in error_msg for keyword in network_keywords)


def simple_network_retry(max_retries: int = 5, base_delay: float = 2.0):
    """简化的网络重试装饰器"""
    config = SimpleNetworkRetryConfig(
        max_retries=max_retries,
        base_delay=base_delay
    )
    
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            last_exception = None
            
            for attempt in range(config.max_retries + 1):
                try:
                    return func(*args, **kwargs)
                
                except Exception as e:
                    last_exception = e
                    
                    if attempt == config.max_retries:
                        print(f"❌ 网络操作最终失败，已重试 {config.max_retries} 次")
                        raise e
                    
                    if not is_network_error(e):
                        print(f"❌ 非网络错误，不进行重试: {str(e)}")
                        raise e
                    
                    delay = calculate_delay(attempt, config)
                    print(f"⚠️  网络错误 (尝试 {attempt + 1}/{config.max_retries + 1}): {str(e)}")
                    print(f"⏳ 等待 {delay:.1f} 秒后重试...")
                    
                    time.sleep(delay)
            
            raise last_exception
        
        return wrapper
    return decorator


def test_basic_connectivity() -> Dict[str, Any]:
    """测试基本网络连接"""
    print("🔍 测试网络连接...")
    
    test_urls = [
        "https://www.google.com",
        "https://huggingface.co", 
        "https://github.com"
    ]
    
    results = {
        "connected": False,
        "successful_urls": [],
        "failed_urls": []
    }
    
    for url in test_urls:
        try:
            response = requests.get(url, timeout=10)
            if response.status_code == 200:
                results["successful_urls"].append(url)
                results["connected"] = True
                print(f"   ✅ {url}")
            else:
                results["failed_urls"].append((url, f"HTTP {response.status_code}"))
                print(f"   ❌ {url}: HTTP {response.status_code}")
        
        except Exception as e:
            results["failed_urls"].append((url, str(e)))
            print(f"   ❌ {url}: {e}")
    
    return results


def get_simple_recommendations(connectivity_results: Dict[str, Any]) -> List[str]:
    """根据连接测试结果提供简单建议"""
    recommendations = []
    
    if not connectivity_results["connected"]:
        recommendations.extend([
            "❌ 网络连接异常，建议:",
            "   1. 检查网络连接",
            "   2. 尝试关闭代理服务器", 
            "   3. 重启网络设备",
            "   4. 检查防火墙设置",
            "   5. 联系网络管理员"
        ])
    else:
        success_count = len(connectivity_results["successful_urls"])
        total_count = success_count + len(connectivity_results["failed_urls"])
        
        if success_count == total_count:
            recommendations.append("✅ 网络连接正常")
        else:
            recommendations.append("⚠️  部分网站无法访问，可能影响上传")
        
        recommendations.extend([
            "💡 上传优化建议:",
            "   1. 使用稳定的网络环境",
            "   2. 避免在网络繁忙时上传",
            "   3. 考虑减小文件分片大小"
        ])
    
    return recommendations


@simple_network_retry(max_retries=5, base_delay=2.0)
def robust_hub_upload(
    dataset,
    repo_id: str,
    max_shard_size: str = "200MB",
    private: bool = False
):
    """带重试机制的 Hub 上传"""
    print(f"🚀 开始上传数据集到 {repo_id}...")
    
    return dataset.push_to_hub(
        repo_id=repo_id,
        max_shard_size=max_shard_size,
        private=private
    )


def upload_with_network_check(
    dataset,
    repo_id: str,
    max_shard_size: str = "200MB", 
    private: bool = False,
    skip_network_check: bool = False
):
    """带网络检查的上传函数"""
    
    if not skip_network_check:
        print("🔍 检查网络连接...")
        connectivity = test_basic_connectivity()
        recommendations = get_simple_recommendations(connectivity)
        
        for rec in recommendations:
            print(f"   {rec}")
        
        if not connectivity["connected"]:
            print("❌ 网络连接检查失败，建议解决网络问题后再试")
            raise ConnectionError("网络连接不可用")
    
    try:
        start_time = time.time()
        
        result = robust_hub_upload(
            dataset=dataset,
            repo_id=repo_id,
            max_shard_size=max_shard_size,
            private=private
        )
        
        duration = time.time() - start_time
        print(f"✅ 上传成功! 耗时: {duration:.1f} 秒")
        
        return result
    
    except Exception as e:
        duration = time.time() - start_time
        print(f"❌ 上传失败! 耗时: {duration:.1f} 秒")
        
        # 提供详细的错误建议
        error_msg = str(e).lower()
        
        if "proxy" in error_msg:
            print("\n🔧 代理问题解决建议:")
            print("   - 尝试关闭代理服务器")
            print("   - 检查代理设置是否正确")
            print("   - 使用直连网络")
        
        elif "connection reset" in error_msg or "connection broken" in error_msg:
            print("\n🔧 连接问题解决建议:")
            print("   - 检查网络稳定性")
            print("   - 减小分片大小 (如 MAX_SHARD_SIZE = '50MB')")
            print("   - 在网络较好的时间重试")
        
        elif "timeout" in error_msg:
            print("\n🔧 超时问题解决建议:")
            print("   - 增加网络超时时间")
            print("   - 检查网络速度")
            print("   - 尝试分批上传")
        
        elif "authentication" in error_msg or "permission" in error_msg:
            print("\n🔧 认证问题解决建议:")
            print("   - 重新登录: huggingface-cli login")
            print("   - 检查token权限")
            print("   - 确认仓库名格式正确")
        
        else:
            print("\n🔧 通用解决建议:")
            print("   - 检查网络连接稳定性")
            print("   - 重新登录 Hugging Face")
            print("   - 稍后重试")
        
        raise


def create_simple_network_test():
    """创建简单的网络测试脚本"""
    test_script = '''#!/usr/bin/env python3
# simple_network_test.py - 简单网络测试
import requests
import time

def test_network():
    print("🔍 测试网络连接...")
    
    urls = [
        "https://www.google.com",
        "https://huggingface.co",
        "https://github.com"
    ]
    
    all_good = True
    
    for url in urls:
        try:
            start = time.time()
            response = requests.get(url, timeout=10)
            duration = time.time() - start
            
            if response.status_code == 200:
                print(f"✅ {url}: {duration:.2f}s")
            else:
                print(f"❌ {url}: HTTP {response.status_code}")
                all_good = False
        except Exception as e:
            print(f"❌ {url}: {e}")
            all_good = False
    
    if all_good:
        print("\\n🎉 网络连接正常！")
    else:
        print("\\n⚠️  网络可能有问题，请检查连接")

if __name__ == "__main__":
    test_network()
'''
    
    try:
        with open("simple_network_test.py", "w", encoding="utf-8") as f:
            f.write(test_script)
        print("📝 已创建网络测试脚本: simple_network_test.py")
        return True
    except Exception as e:
        print(f"❌ 创建测试脚本失败: {e}")
        return False


# 使用示例
def demo_usage():
    """使用示例"""
    print("=" * 50)
    print("简化网络工具使用示例")
    print("=" * 50)
    
    # 测试网络连接
    connectivity = test_basic_connectivity()
    recommendations = get_simple_recommendations(connectivity)
    
    print("\n💡 建议:")
    for rec in recommendations:
        print(rec)
    
    # 创建测试脚本
    create_simple_network_test()


if __name__ == "__main__":
    demo_usage()